import cv2
import SimpleITK as sitk
import matplotlib.pyplot as plt
import numpy as np
from skimage import feature, morphology

"""
NOTES:
    1. 滤波可能降低分割质量，薄的部分会消失
    2. 考虑先旋转？
"""


def test():
    image = sitk.ReadImage(r'C:\study\FFR\myProject\mytest\mydata_apply\input\case_00217\segmentation.nii.gz')
    array = sitk.GetArrayFromImage(image)
    # array = array.swapaxes(0, 2)
    test_slice = 85
    fig = plt.figure()
    axis = fig.add_subplot(1, 4, 1)
    axis.imshow(array[test_slice])

    # output = cv2.blur(array[test_slice], ksize=(15, 15), borderType=cv2.BORDER_REPLICATE)
    # axis = fig.add_subplot(1, 3, 2)
    # axis.imshow(output)

    array = array.astype(np.uint8)
    output2 = cv2.medianBlur(array[test_slice], ksize=9)
    axis = fig.add_subplot(1, 4, 2)
    axis.imshow(output2)

    skele = morphology.skeletonize(array[test_slice])
    axis = fig.add_subplot(1, 4, 3)
    axis.imshow(skele)

    output4 = np.logical_or(output2, skele)
    output4 = output4.astype(np.float32)
    axis = fig.add_subplot(1, 4, 4)
    axis.imshow(output4)

    plt.show()


def get_closing(array, radius=15, out=None):
    """
    calculate closing along Z, Y and X axis

    :param array:
    :param radius: radius of disk footprint
    :param out: output reference
    :return:
    """
    array = array.astype(np.uint8)
    k = morphology.disk(radius)
    if out is None:
        ret = array.copy()
    else:
        ret = out

    for z in range(array.shape[0]):
        ret[z, :, :] = morphology.closing(ret[z, :, :], k)
    print('------- finished Z-Axis closing ---------')
    for y in range(array.shape[1]):
        ret[:, y, :] = morphology.closing(ret[:, y, :], k)
    print('------- finished Y-Axis closing ---------')
    for x in range(array.shape[2]):
        ret[:, :, x] = morphology.closing(ret[:, :, x], k)
    print('------- finished X-Axis closing ---------')
    return ret


def get_blur(array, ksize=9, out=None):
    """
    calculate median blur

    :param array:
    :param ksize: width of kernel window
    :param out: output reference
    :return:
    """
    array = array.astype(np.uint8)
    if out is None:
        ret = array.copy()
    else:
        ret = out

    for z in range(array.shape[0]):
        ret[z, :, :] = cv2.medianBlur(ret[z, :, :], ksize)
    for y in range(array.shape[1]):
        ret[:, y, :] = cv2.medianBlur(ret[:, y, :], ksize)
    for x in range(array.shape[2]):
        ret[:, :, x] = cv2.medianBlur(ret[:, :, x], ksize)

    return ret


def antialias(in_path, out_path, radius, ksize):
    """
    将in_path的CTA图像闭运算、平滑化后，存储到out_path

    :param in_path:
    :param out_path:
    :return:
    """
    image = sitk.ReadImage(in_path)
    array = sitk.GetArrayFromImage(image)
    array = array.astype(np.uint8)

    get_closing(array, radius=radius, out=array)
    get_blur(array, ksize=ksize, out=array)

    out_image = sitk.GetImageFromArray(array)
    out_image.CopyInformation(image)
    sitk.WriteImage(out_image, out_path)


def __common_filter(array, i, j, ksize):
    """
    return the most commonly appeared non-zero label value

    :param array:
    :param i:
    :param j:
    :param ksize:
    :return:
    """
    radius = ksize // 2
    x1 = i - radius
    y1 = j - radius
    x2 = i + radius
    y2 = j + radius
    cnt_dic = dict({})
    for x in range(x1, x2 + 1):
        if x < 0 or x >= array.shape[0]:
            continue
        for y in range(y1, y2 + 1):
            if y < 0 or y >= array.shape[1]:
                continue
            key = array[x, y]
            if key in cnt_dic:
                cnt_dic[key] += 1
            else:
                cnt_dic[key] = 1
    max_cnt = -1
    max_val = 0
    second_cnt = -1
    second_val = 0



    for k, v in cnt_dic.items():
        if v > max_cnt or (v == max_cnt and k > max_val):
            second_cnt = max_cnt
            second_val = max_val
            max_cnt = v
            max_val = k
        elif v > second_cnt or (v == second_cnt and k > second_val):
            second_cnt = v
            second_val = k
    if max_val != 0:
        return max_val
    else:
        return second_val


def commonFilter2D(array, ksize, out=None):
    """
    return the filtered array that turns each non-zero pixel to the most existed value around it
    :param array:
    :param ksize: must be an odd number
    :return:
    """

    ret = out
    if ret is None:
        ret = np.zeros(array.shape)
    ret = ret.astype(np.uint8)
    for i in range(array.shape[0]):
        for j in range(array.shape[1]):
            if array[i, j] != 0:
                ret[i, j] = __common_filter(array, i, j, ksize)
    return ret

if __name__ == '__main__':
    in_path = r'C:\study\FFR\myProject\mytest\mydata_apply\input\case_00144\segmentation.nii.gz'
    out_path = r'temp_test_antialias_00144.nii.gz'
    # antialias(in_path, out_path)
    # test()

    # test denoising the edge of sectional labels
    in_path = 'C:/study/FFR/myProject/mytest/mydata_test/pred_recover/results_recovered/prediction_00999_noised.nii.gz'
    out_path = 'C:/study/FFR/myProject/mytest/mydata_test/pred_recover/results_recovered/prediction_00999.nii.gz'
    image = sitk.ReadImage(in_path)
    array = sitk.GetArrayFromImage(image)
    array = array.astype(np.uint8)

    ksize = 15
    iter_num = 2

    for i in range(iter_num):
        for z in range(array.shape[0]):
            array[z, :, :] = commonFilter2D(array[z, :, :], ksize)
        print('finished z-axis')
        for y in range(array.shape[1]):
            array[:, y, :] = commonFilter2D(array[:, y, :], ksize)
        print('finished y-axis')
        for x in range(array.shape[2]):
            array[:, :, x] = commonFilter2D(array[:, :, x], ksize)
        print('finished x-axis')

        print('finished', i, 'iteration')

    out_img = sitk.GetImageFromArray(array)
    out_img.CopyInformation(image)
    sitk.WriteImage(out_img, out_path)
